package org.dromara.common.mybatis.interceptor

import com.baomidou.mybatisplus.core.plugins.InterceptorIgnoreHelper
import com.baomidou.mybatisplus.core.toolkit.PluginUtils
import com.baomidou.mybatisplus.extension.plugins.handler.MultiDataPermissionHandler
import com.baomidou.mybatisplus.extension.plugins.inner.BaseMultiTableInnerInterceptor
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor
import io.github.oshai.kotlinlogging.KotlinLogging
import net.sf.jsqlparser.expression.Expression
import net.sf.jsqlparser.schema.Table
import net.sf.jsqlparser.statement.delete.Delete
import net.sf.jsqlparser.statement.select.PlainSelect
import net.sf.jsqlparser.statement.select.Select
import net.sf.jsqlparser.statement.select.SetOperationList
import net.sf.jsqlparser.statement.update.Update
import org.apache.ibatis.executor.Executor
import org.apache.ibatis.executor.statement.StatementHandler
import org.apache.ibatis.mapping.BoundSql
import org.apache.ibatis.mapping.MappedStatement
import org.apache.ibatis.mapping.SqlCommandType
import org.apache.ibatis.session.ResultHandler
import org.apache.ibatis.session.RowBounds
import org.dromara.common.mybatis.handler.PlusDataPermissionHandler
import java.sql.Connection
import java.sql.SQLException
import java.util.function.Consumer

/**
 * 数据权限拦截器
 *
 * @author Lion Li
 * @version 3.5.0
 * @updater LikeYouDo
 * @date 2025/1/27 3:47
 */
open class PlusDataPermissionInterceptor() : BaseMultiTableInnerInterceptor(), InnerInterceptor {

    private val dataPermissionHandler: PlusDataPermissionHandler = PlusDataPermissionHandler()

    companion object{
        private val log = KotlinLogging.logger {  }
    }

    /**
     * 在执行查询之前，检查并处理数据权限相关逻辑
     *
     * @param executor      MyBatis 执行器对象
     * @param ms            映射语句对象
     * @param parameter     方法参数
     * @param rowBounds     分页对象
     * @param resultHandler 结果处理器
     * @param boundSql      绑定的 SQL 对象
     * @throws SQLException 如果发生 SQL 异常
     */
    @Throws(SQLException::class)
    override fun beforeQuery(
        executor: Executor?,
        ms: MappedStatement,
        parameter: Any?,
        rowBounds: RowBounds?,
        resultHandler: ResultHandler<*>?,
        boundSql: BoundSql?
    ) {
        // 检查是否需要忽略数据权限处理
        if (InterceptorIgnoreHelper.willIgnoreDataPermission(ms.id)) {
            return
        }
        // 检查是否缺少有效的数据权限注解
        if (dataPermissionHandler.invalid()) {
            return
        }
        // 解析 sql 分配对应方法
        val mpBs = PluginUtils.mpBoundSql(boundSql)
        mpBs.sql(parserSingle(mpBs.sql(), ms.id))
    }

    /**
     * 在准备 SQL 语句之前，检查并处理更新和删除操作的数据权限相关逻辑
     *
     * @param sh                 MyBatis StatementHandler 对象
     * @param connection         数据库连接对象
     * @param transactionTimeout 事务超时时间
     */
    override fun beforePrepare(sh: StatementHandler?, connection: Connection?, transactionTimeout: Int?) {
        val mpSh = PluginUtils.mpStatementHandler(sh)
        val ms = mpSh.mappedStatement()
        // 获取 SQL 命令类型（增、删、改、查）
        val sct = ms.sqlCommandType

        // 只处理更新和删除操作的 SQL 语句
        if (sct == SqlCommandType.UPDATE || sct == SqlCommandType.DELETE) {
            if (InterceptorIgnoreHelper.willIgnoreDataPermission(ms.id)) {
                return
            }
            // 检查是否缺少有效的数据权限注解
            if (dataPermissionHandler.invalid()) {
                return
            }
            val mpBs = mpSh.mPBoundSql()
            mpBs.sql(parserMulti(mpBs.sql(), ms.id))
        }
    }

    /**
     * 处理 SELECT 查询语句中的 WHERE 条件
     *
     * @param select SELECT 查询对象
     * @param index  查询语句的索引
     * @param sql    查询语句
     * @param obj    WHERE 条件参数
     */
    override fun processSelect(select: Select?, index: Int, sql: String?, obj: Any?) {
        if (select is PlainSelect) {
            this.setWhere(select, obj as String)
        } else if (select is SetOperationList) {
            val selectBodyList = select.selects
            selectBodyList.forEach(Consumer { this.setWhere(it as PlainSelect, obj as String) })
        }
    }

    /**
     * 处理 UPDATE 语句中的 WHERE 条件
     *
     * @param update UPDATE 查询对象
     * @param index  查询语句的索引
     * @param sql    查询语句
     * @param obj    WHERE 条件参数
     */
    override fun processUpdate(update: Update, index: Int, sql: String?, obj: Any?) {
        val sqlSegment = dataPermissionHandler.getSqlSegment(
            update.where,
            false
        )
        if (null != sqlSegment) {
            update.where = sqlSegment
        }
    }

    /**
     * 处理 DELETE 语句中的 WHERE 条件
     *
     * @param delete DELETE 查询对象
     * @param index  查询语句的索引
     * @param sql    查询语句
     * @param obj    WHERE 条件参数
     */
    override fun processDelete(delete: Delete, index: Int, sql: String?, obj: Any?) {
        val sqlSegment = dataPermissionHandler.getSqlSegment(
            delete.where,
            false
        )
        if (null != sqlSegment) {
            delete.where = sqlSegment
        }
    }

    /**
     * 设置 SELECT 语句的 WHERE 条件
     *
     * @param plainSelect       SELECT 查询对象
     * @param mappedStatementId 映射语句的 ID
     */
    protected fun setWhere(plainSelect: PlainSelect, mappedStatementId: String) {
        val sqlSegment = dataPermissionHandler.getSqlSegment(
            plainSelect.where,
            true
        )
        if (null != sqlSegment) {
            plainSelect.where = sqlSegment
        }
    }

    /**
     * 构建表达式，用于处理表的数据权限
     *
     * @param table        表对象
     * @param where        WHERE 条件表达式
     * @param whereSegment WHERE 条件片段
     * @return 构建的表达式
     */
    override fun buildTableExpression(table: Table?, where: Expression?, whereSegment: String?): Expression {
        // 只有新版数据权限处理器才会执行到这里
        val handler = dataPermissionHandler as MultiDataPermissionHandler
        return handler.getSqlSegment(table, where, whereSegment)
    }
}
